from py2neo import Graph
import question_classifier
import model_process
import pickle


class Predict:
    def __init__(self, tg):
        self.tg = tg

    def run(self, query, question):
        if question == 0:
            return self.director(query)
        elif question == 1:
            return self.screenwriter(query)
        elif question == 2:
            return self.actor(query)
        elif question == 3:
            return self.type(query)
        elif question == 4:
            return self.area(query)
        elif question == 5:
            return self.language(query)
        elif question == 6:
            return self.director_screenwriter(query)
        elif question == 7:
            return self.director_actor(query)
        elif question == 8:
            return self.director_type(query)
        elif question == 9:
            return self.director_area(query)
        elif question == 10:
            return self.screenwriter_type(query)
        elif question == 11:
            return self.actor_type(query)
        elif question == 12:
            return self.actor_language(query)

    def director(self, query):
        query = query.split(" ")
        result = self.tg.run("match(n:Movie)-[r:direct]->(b:Director) where n.m_name=$m_name return b.dir_name",
                             m_name=query[0]).to_data_frame()
        if not result.empty:
            # print(query[0] + "的导演是：" + result.iat[0, 0])
            return query[0] + "的导演是：" + result.iat[0, 0]
        else:
            return "很抱歉，这个问题我还不会回答，或者您可以换一种问法"

    def screenwriter(self, query):
        query = query.split(" ")
        result = self.tg.run("match(n:Movie)-[r:write]->(b:Screenwriter) where n.m_name=$m_name return b.sw_name",
                             m_name=query[0]).to_data_frame()
        if not result.empty:
            word = ""
            for i in range(len(result) - 1):
                word += result.iat[i, 0] + "、"
            word += result.iat[len(result) - 1, 0]
            return query[0] + "的编剧是：" + word
        else:
            return "很抱歉，这个问题我还不会回答，或者您可以换一种问法"

    def actor(self, query):
        query = query.split(" ")
        result = self.tg.run("match(n:Movie)-[r:act_in]->(b:Actor) where n.m_name=$m_name return b.act_name",
                             m_name=query[0]).to_data_frame()
        if not result.empty:
            word = ""
            for i in range(len(result) - 1):
                word += result.iat[i, 0] + "、"
            word += result.iat[len(result) - 1, 0]
            # print(query[0] + "的主要演员有：" + word)
            return query[0] + "的主要演员有：" + word
        else:
            return "很抱歉，这个问题我还不会回答，或者您可以换一种问法"

    def type(self, query):
        query = query.split(" ")
        result = self.tg.run("match(n:Movie)-[r:is]->(b:Type) where n.m_name=$m_name return b.type",
                             m_name=query[0]).to_data_frame()
        if not result.empty:
            word = ""
            for i in range(len(result) - 1):
                word += result.iat[i, 0] + "、"
            word += result.iat[len(result) - 1, 0]
            return query[0] + "的类型是：" + word
        else:
            return "很抱歉，这个问题我还不会回答，或者您可以换一种问法"

    def area(self, query):
        query = query.split(" ")
        result = self.tg.run("match(n:Movie)-[r:release_in]->(b:Area) where n.m_name=$m_name return b.area",
                             m_name=query[0]).to_data_frame()
        if not result.empty:
            return query[0] + "的上映地区是：" + result.iat[0, 0]
        else:
            return "很抱歉，这个问题我还不会回答，或者您可以换一种问法"

    def language(self, query):
        query = query.split(" ")
        result = self.tg.run("match(n:Movie)-[r:language_in]->(b:Language) where n.m_name=$m_name return b.language",
                             m_name=query[0]).to_data_frame()
        if not result.empty:
            word = ""
            for i in range(len(result) - 1):
                word += result.iat[i, 0] + "、"
            word += result.iat[len(result) - 1, 0]
            return query[0] + "的语言有：" + word
        else:
            return "很抱歉，这个问题我还不会回答，或者您可以换一种问法"

    def director_screenwriter(self, query):
        query = query.split(" ")
        result = self.tg.run("match(m:Movie)-[:direct]-(d:Director) where d.dir_name=$dir_name " +
                             "match(m)-[:write]-(s:Screenwriter) where s.sw_name=$sw_name return m.m_name",
                             dir_name=query[0], sw_name=query[1]).to_data_frame()
        if not result.empty:
            word = ""
            for i in range(len(result) - 1):
                word += result.iat[i, 0] + "、"
            word += result.iat[len(result) - 1, 0]
            # print(query[0] + "和" + query[1] + "合作过的电影有" + word)
            return query[0] + "和" + query[1] + "合作过的电影有：" + word
        else:
            return "很抱歉，这个问题我还不会回答，或者您可以换一种问法"

    def director_actor(self, query):
        query = query.split(" ")
        result = self.tg.run("match(m:Movie)-[:direct]-(d:Director) where d.dir_name=$dir_name " +
                             "match(m)-[:act_in]-(a:Actor) where a.act_name=$act_name return m.m_name",
                             dir_name=query[0], act_name=query[1]).to_data_frame()
        if not result.empty:
            word = ""
            for i in range(len(result) - 1):
                word += result.iat[i, 0] + "、"
            word += result.iat[len(result) - 1, 0]
            # print(query[0] + "和" + query[1] + "合作过的电影有" + word)
            return query[0] + "和" + query[1] + "合作过的电影有：" + word
        else:
            return "很抱歉，这个问题我还不会回答，或者您可以换一种问法"

    def director_type(self, query):
        query = query.split(" ")
        result = self.tg.run("match(n:Director)-[r:dir_type]->(b:Type) where n.dir_name=$dir_name return b.type",
                             dir_name=query[0]).to_data_frame()
        if not result.empty:
            word = ""
            for i in range(len(result) - 1):
                word += result.iat[i, 0] + "、"
            word += result.iat[len(result) - 1, 0]
            return query[0] + "导演拍过的电影类型有：" + word
        else:
            return "很抱歉，这个问题我还不会回答，或者您可以换一种问法"

    def director_area(self, query):
        query = query.split(" ")
        result = self.tg.run("match(n:Director)-[r:dir_in]->(b:Area) where n.dir_name=$dir_name return b.area",
                             dir_name=query[0]).to_data_frame()
        if not result.empty:
            word = ""
            for i in range(len(result) - 1):
                word += result.iat[i, 0] + "、"
            word += result.iat[len(result) - 1, 0]
            return query[0] + "出品电影的地区有：" + word
        else:
            return "很抱歉，这个问题我还不会回答，或者您可以换一种问法"

    def screenwriter_type(self, query):
        query = query.split(" ")
        result = self.tg.run("match(n:Screenwriter)-[r:write_type]->(b:Type) where n.sw_name=$sw_name return b.type",
                             sw_name=query[0]).to_data_frame()
        if not result.empty:
            word = ""
            for i in range(len(result) - 1):
                word += result.iat[i, 0] + "、"
            word += result.iat[len(result) - 1, 0]
            return query[0] + "编剧写过的电影类型有：" + word
        else:
            return "很抱歉，这个问题我还不会回答，或者您可以换一种问法"

    def actor_type(self, query):
        query = query.split(" ")
        result = self.tg.run("match(n:Actor)-[r:act_type]->(b:Type) where n.act_name=$act_name return b.type",
                             act_name=query[0]).to_data_frame()
        if not result.empty:
            word = ""
            for i in range(len(result) - 1):
                word += result.iat[i, 0] + "、"
            word += result.iat[len(result) - 1, 0]
            return query[0] + "出演过的电影类型有：" + word
        else:
            return "很抱歉，这个问题我还不会回答，或者您可以换一种问法"

    def actor_language(self, query):
        query = query.split(" ")
        result = self.tg.run("match(n:Actor)-[r:speak]->(b:Language) where n.act_name=$act_name return b.language",
                             act_name=query[0]).to_data_frame()
        if not result.empty:
            word = ""
            for i in range(len(result) - 1):
                word += result.iat[i, 0] + "、"
            word += result.iat[len(result) - 1, 0]
            return query[0] + "会说的语言有：" + word
        else:
            return "很抱歉，这个问题我还不会回答，或者您可以换一种问法"


test_graph = Graph(
    host="127.0.0.1",
    http_port=7474,
    user="neo4j",
    password="neo4j"
)


if __name__ == "__main__":
    data_dir = './newquestion/'
    pkl_file = open("./newdict/vocabulary.pkl", "rb")
    vocabulary = pickle.load(pkl_file)
    sentence = "陈凯歌在哪个地区"
    model = model_process.NaiveBayesModelMe()
    # model.fit()
    prediction = model.test(sentence, vocabulary)
    query = question_classifier.match_question(prediction, sentence)

    predict = Predict(test_graph)
    predict.run(query, prediction)
